package org.backend.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

import static org.backend.utils.JwtUtil.getOpenId;

@Component
@Slf4j
public class WxWebSocketHandler extends TextWebSocketHandler {

    private String openid;

    // 用来保存所有连接的 WebSocket 会话
    private static final CopyOnWriteArraySet<WebSocketSession> webSocketSet = new CopyOnWriteArraySet<>();

    private static final Map<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String token = session.getHandshakeHeaders().getFirst("Authorization");
        if (token == null) {
            throw new IllegalArgumentException("Authorization header missing");
        }
        openid = getOpenId(token);
        if (openid != null) {
            userSessions.put(openid, session);
            log.info("WebSocket connection established for user: {}", openid);
        } else {
            log.info("Failed to retrieve openid from Authorization header");
        }

    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 处理从前端接收到的消息 （如用户请求）
        log.info("用户{}发送消息: {}", session.getId(), message.getPayload());
        session.sendMessage(new TextMessage(message.getPayload()));
    }

    public void sendNotificationToUser(String openid, String message, Object data) {
        WebSocketSession session = userSessions.get(openid);
        if (session != null && session.isOpen()) {
            try {
                if (data == null) {
                    session.sendMessage(new TextMessage(message));
                    log.info("Message sent to user {}: {}", openid, message);
                } else {
                    Map<String, Object> response = new HashMap<>();
                    response.put("data", data);
                    response.put("message", message);

                    ObjectMapper objectMapper = new ObjectMapper();
                    String jsonResponse = objectMapper.writeValueAsString(response);
                    log.info("\n\njsonResponse: {}\n\n", jsonResponse);
                    session.sendMessage(new TextMessage(jsonResponse));
                    log.info("Message with data sent to user {}: {}", openid, jsonResponse);
                }

            } catch (IOException e) {
                log.error("Error sending message to user {}: {}", openid, e.getMessage());
            }
        }
    }

    // 获取所有连接的 openid 列表
    public Set<String> getAllOpenids() {
        return userSessions.keySet();
    }

    // 获取用户的 WebSocket 会话
    public WebSocketSession getSessionByOpenId(String openid) {
        return userSessions.get(openid);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        if (openid != null) {
            userSessions.remove(openid);
            log.info("WebSocket connection closed for user: {}", openid);
        }
    }

}
